import os
import torch
import torchvision
import torchvision.transforms as transforms
from torch import nn
from torch.utils.data import DataLoader
import numpy as np
import matplotlib.pyplot as plt
from torchvision.utils import save_image
import time


# ========================= 创建目录结构 =========================
def create_directories():
    dirs = [
        "original_model",      # 原始模型文件夹
        "robust_model",        # 增强模型文件夹
        "adversarial_examples"  # 对抗样本示例文件夹
    ]
    for d in dirs:
        if not os.path.exists(d):
            os.makedirs(d)
            print(f"创建目录: {d}")


# ========================= 数据加载与预处理 =========================
def load_data(batch_size=64):
    transform = transforms.ToTensor()
    train_data = torchvision.datasets.MNIST(root='data', train=True, download=True, transform=transform)
    test_data = torchvision.datasets.MNIST(root='data', train=False, download=True, transform=transform)
    train_dataloader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True)
    test_dataloader = DataLoader(dataset=test_data, batch_size=batch_size)
    return train_dataloader, test_dataloader


# ========================= 模型定义 =========================
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 6, 3, stride=1, padding=1),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5, stride=1, padding=1),
            nn.MaxPool2d(2, 2)
        )
        self.fc = nn.Sequential(
            nn.Linear(576, 120),
            nn.Linear(120, 84),
            nn.Linear(84, 10)
        )

    def forward(self, x):
        out = self.conv(x)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out


# ========================= 常规训练函数 =========================
def train_model(network, train_dataloader, device, loss_fn, optimizer, epochs=10, 
               model_path="model.pth", loss_plot_path="loss.png"):
    network.train()
    losses = []
    start_time = time.time()
    
    for epoch in range(epochs):
        epoch_loss = 0.0
        for X, y in train_dataloader:
            X, y = X.to(device), y.to(device)
            
            # 前向传播
            pred = network(X)
            loss = loss_fn(pred, y)
            
            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            epoch_loss += loss.item() * X.size(0)
            
        # 计算平均损失
        avg_loss = epoch_loss / len(train_dataloader.dataset)
        losses.append(avg_loss)
        
        # 打印训练信息
        elapsed = time.time() - start_time
        print(f"Epoch {epoch+1}/{epochs} - Loss: {avg_loss:.6f} - Time: {elapsed:.2f}s")
    
    # 保存模型
    torch.save(network.state_dict(), model_path)
    print(f"模型已保存到 {model_path}")
    
    # 绘制损失曲线
    plt.figure(figsize=(10, 6))
    plt.plot(range(1, epochs+1), losses, 'o-')
    plt.title("Training Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.grid(True)
    plt.savefig(loss_plot_path)
    plt.close()
    print(f"损失曲线已保存到 {loss_plot_path}")
    
    return losses


# ========================= 在线对抗训练函数 =========================
def adversarial_train(network, train_dataloader, device, loss_fn, optimizer, 
                     epochs=10, epsilon=0.1, alpha=0.5, 
                     model_path="robust_model.pth", loss_plot_path="robust_loss.png"):
    """
    在线对抗训练函数，在每个批次中生成对抗样本
    
    参数:
        network: 待训练的网络模型
        train_dataloader: 训练数据加载器
        device: 计算设备（CPU/GPU）
        loss_fn: 损失函数
        optimizer: 优化器
        epochs: 训练轮数
        epsilon: FGSM扰动大小
        alpha: 原始样本和对抗样本损失的权重比例（alpha用于原始样本损失）
        model_path: 模型保存路径
        loss_plot_path: 损失曲线保存路径
    """
    network.train()
    losses = []
    clean_losses = []
    adv_losses = []
    start_time = time.time()    
    # 保存一些示例对抗样本的计数器
    sample_counter = 0    
    for epoch in range(epochs):
        epoch_loss = 0.0
        epoch_clean_loss = 0.0
        epoch_adv_loss = 0.0
        batch_count = 0
        
        for X, y in train_dataloader:
            batch_count += 1
            X, y = X.to(device), y.to(device)
            batch_size = X.size(0)            
            # 计算原始样本的损失
            X.requires_grad = True
            outputs_clean = network(X)
            loss_clean = loss_fn(outputs_clean, y)            
            # 生成对抗样本
            network.zero_grad()
            loss_clean.backward(retain_graph=True)
            X_grad = X.grad.data            
            # FGSM攻击
            X_adv = X.detach() + epsilon * X_grad.sign()
            X_adv = torch.clamp(X_adv, 0, 1)            
            #  计算对抗样本的损失
            outputs_adv = network(X_adv)
            loss_adv = loss_fn(outputs_adv, y)            
            # 合并损失
            loss = alpha * loss_clean + (1 - alpha) * loss_adv            
            # 优化模型参数
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            # 累计损失
            epoch_loss += loss.item() * batch_size
            epoch_clean_loss += loss_clean.item() * batch_size
            epoch_adv_loss += loss_adv.item() * batch_size
            
            # 保存一些对抗样本示例
            if sample_counter < 5 and epoch == 0 and batch_count % 100 == 0:
                examples_dir = "adversarial_examples"
                # 保存原始图像
                original_path = os.path.join(examples_dir, f"online_original_{sample_counter}.png")
                save_image(X.detach().cpu()[:8], original_path, nrow=4)
                
                # 保存对抗样本
                adv_path = os.path.join(examples_dir, f"online_adversarial_{sample_counter}.png")
                save_image(X_adv.detach().cpu()[:8], adv_path, nrow=4)
                
                # 保存扰动（放大5倍以便观察）
                perturbation = (X_adv - X).detach().cpu() * 5
                perturbation = torch.clamp(perturbation + 0.5, 0, 1)
                pert_path = os.path.join(examples_dir, f"online_perturbation_{sample_counter}.png")
                save_image(perturbation[:8], pert_path, nrow=4)
                
                sample_counter += 1
        
        # 计算平均损失
        avg_loss = epoch_loss / len(train_dataloader.dataset)
        avg_clean_loss = epoch_clean_loss / len(train_dataloader.dataset)
        avg_adv_loss = epoch_adv_loss / len(train_dataloader.dataset)
        
        losses.append(avg_loss)
        clean_losses.append(avg_clean_loss)
        adv_losses.append(avg_adv_loss)
        
        # 打印训练信息
        elapsed = time.time() - start_time
        print(f"Epoch {epoch+1}/{epochs} - Total Loss: {avg_loss:.6f} - Clean Loss: {avg_clean_loss:.6f} - Adv Loss: {avg_adv_loss:.6f} - Time: {elapsed:.2f}s")
    
    # 保存模型
    torch.save(network.state_dict(), model_path)
    print(f"增强模型已保存到 {model_path}")
    
    # 绘制损失曲线
    plt.figure(figsize=(10, 6))
    plt.plot(range(1, epochs+1), losses, 'o-', label='Total Loss')
    plt.plot(range(1, epochs+1), clean_losses, 's-', label='Clean Loss')
    plt.plot(range(1, epochs+1), adv_losses, '^-', label='Adversarial Loss')
    plt.title("Adversarial Training Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.legend()
    plt.grid(True)
    plt.savefig(loss_plot_path)
    plt.close()
    print(f"损失曲线已保存到 {loss_plot_path}")
    
    return losses


# ========================= 测试函数 =========================
def evaluate_model(network, test_dataloader, device, description=""):
    network.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for X, y in test_dataloader:
            X, y = X.to(device), y.to(device)
            outputs = network(X)
            _, predicted = torch.max(outputs.data, 1)
            total += y.size(0)
            correct += (predicted == y).sum().item()
    
    accuracy = correct / total
    print(f"{description} 准确率: {accuracy * 100:.2f}%")
    return accuracy


# ========================= FGSM攻击测试 =========================
def test_fgsm_attack(network, test_dataloader, device, loss_fn, epsilons, 
                    output_dir, plot_filename="accuracy_vs_epsilon.png"):
    network.eval()
    accuracies = []
    
    # 为每个epsilon值计算准确率
    for epsilon in epsilons:
        correct = 0
        total = 0
        
        for X, y in test_dataloader:
            X, y = X.to(device), y.to(device)
            X.requires_grad = True
            
            # 前向传播
            outputs = network(X)
            loss = loss_fn(outputs, y)
            
            # 计算梯度
            network.zero_grad()
            loss.backward()
            
            # 创建对抗样本
            X_adv = X + epsilon * X.grad.sign()
            X_adv = torch.clamp(X_adv, 0, 1)
            
            # 在对抗样本上测试
            with torch.no_grad():
                outputs_adv = network(X_adv)
                _, predicted = torch.max(outputs_adv.data, 1)
                total += y.size(0)
                correct += (predicted == y).sum().item()
                
        # 计算当前epsilon的准确率
        accuracy = correct / total
        accuracies.append(accuracy)
        print(f"Epsilon: {epsilon:.2f}, 准确率: {accuracy * 100:.2f}%")
        
        # 在第一个epsilon值处保存示例
        if epsilon == epsilons[0]:
            examples_dir = output_dir
            # 保存原始图像
            original_path = os.path.join(examples_dir, "test_original.png")
            save_image(X.detach().cpu()[:8], original_path, nrow=4)
            
            # 保存对抗样本
            adv_path = os.path.join(examples_dir, f"test_adversarial_{epsilon:.2f}.png")
            save_image(X_adv.detach().cpu()[:8], adv_path, nrow=4)
            
            # 保存扰动
            perturbation = (X_adv - X).detach().cpu() * 5
            perturbation = torch.clamp(perturbation + 0.5, 0, 1)
            pert_path = os.path.join(examples_dir, f"test_perturbation_{epsilon:.2f}.png")
            save_image(perturbation[:8], pert_path, nrow=4)
    
    # 绘制epsilon与准确率的关系图
    plt.figure(figsize=(10, 6))
    plt.plot(epsilons, accuracies, 'o-')
    plt.title("Accuracy Under FGSM Attack")
    plt.xlabel("Epsilon")
    plt.ylabel("Accuracy")
    plt.ylim(0, 1)
    plt.grid(True)
    plot_path = os.path.join(output_dir, plot_filename)
    plt.savefig(plot_path)
    plt.close()
    print(f"准确率图像已保存到 {plot_path}")
    
    return accuracies


# ========================= 主函数 =========================
def main():
    print("启动在线对抗训练实验...")
    
    # 创建目录结构
    create_directories()
    
    # 设置随机种子以确保可重复性
    torch.manual_seed(42)
    np.random.seed(42)
    
    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    
    # 加载数据
    batch_size = 128  # 增大批次大小以加速训练
    train_dataloader, test_dataloader = load_data(batch_size)
    print(f"训练集大小: {len(train_dataloader.dataset)}, 测试集大小: {len(test_dataloader.dataset)}")
    
    # 损失函数
    loss_fn = nn.CrossEntropyLoss()
    
    # 设置FGSM攻击的epsilon值
    test_epsilons = [0.01, 0.05, 0.09, 0.13, 0.17, 0.2]
    
    # ===== 步骤1: 训练原始模型 =====
    original_model_path = os.path.join("original_model", "model.pth")
    original_loss_plot = os.path.join("original_model", "loss.png")
    
    if os.path.exists(original_model_path):
        # 加载已有的原始模型
        original_model = LeNet().to(device)
        original_model.load_state_dict(torch.load(original_model_path))
        print(f"加载已有的原始模型: {original_model_path}")
    else:
        # 训练原始模型
        print("\n===== 训练原始模型 =====")
        original_model = LeNet().to(device)
        original_optimizer = torch.optim.Adam(original_model.parameters(), lr=0.001, weight_decay=1e-5)
        train_model(original_model, train_dataloader, device, loss_fn, original_optimizer, 
                   epochs=10, model_path=original_model_path, loss_plot_path=original_loss_plot)
    
    # 评估原始模型
    print("\n===== 评估原始模型 =====")
    orig_clean_acc = evaluate_model(original_model, test_dataloader, device, 
                                   description="原始模型在干净测试数据上")
    
    # 测试原始模型在FGSM攻击下的性能
    print("\n===== 原始模型在FGSM攻击下的性能 =====")
    orig_attack_accs = test_fgsm_attack(original_model, test_dataloader, device, loss_fn,
                                      test_epsilons, "original_model")
    
    # ===== 步骤2: 使用在线对抗训练训练增强模型 =====
    robust_model_path = os.path.join("robust_model", "model.pth")
    robust_loss_plot = os.path.join("robust_model", "loss.png")
    
    if os.path.exists(robust_model_path):
        # 加载已有的增强模型
        robust_model = LeNet().to(device)
        robust_model.load_state_dict(torch.load(robust_model_path))
        print(f"加载已有的增强模型: {robust_model_path}")
    else:
        # 训练增强模型
        print("\n===== 使用在线对抗训练训练增强模型 =====")
        robust_model = LeNet().to(device)
        robust_optimizer = torch.optim.Adam(robust_model.parameters(), lr=0.001, weight_decay=1e-5)
        
        # 设置在线对抗训练参数
        train_epsilon = 0.1  # 训练时使用的扰动大小
        alpha = 0.5  # 原始样本和对抗样本损失的权重比例
        
        adversarial_train(robust_model, train_dataloader, device, loss_fn, robust_optimizer, 
                        epochs=15, epsilon=train_epsilon, alpha=alpha,
                        model_path=robust_model_path, loss_plot_path=robust_loss_plot)
    
    # 评估增强模型
    print("\n===== 评估增强模型 =====")
    robust_clean_acc = evaluate_model(robust_model, test_dataloader, device,
                                    description="增强模型在干净测试数据上")
    
    # 测试增强模型在FGSM攻击下的性能
    print("\n===== 增强模型在FGSM攻击下的性能 =====")
    robust_attack_accs = test_fgsm_attack(robust_model, test_dataloader, device, loss_fn,
                                        test_epsilons, "robust_model")
    
    # ===== 步骤3: 生成模型比较图 =====
    plt.figure(figsize=(10, 6))
    plt.plot(test_epsilons, orig_attack_accs, 'bo-', label='Original Model')
    plt.plot(test_epsilons, robust_attack_accs, 'ro-', label='Robust Model')
    plt.title("Model Comparison Under FGSM Attack")
    plt.xlabel("Epsilon")
    plt.ylabel("Accuracy")
    plt.ylim(0, 1)
    plt.legend()
    plt.grid(True)
    comparison_path = "model_comparison.png"
    plt.savefig(comparison_path)
    plt.close()
    print(f"模型对比图已保存到 {comparison_path}")
    
    # ===== 步骤4: 结果汇总 =====
    print("\n===== 实验结果汇总 =====")
    print(f"原始模型在干净测试数据上的准确率: {orig_clean_acc * 100:.2f}%")
    print(f"增强模型在干净测试数据上的准确率: {robust_clean_acc * 100:.2f}%")
    
    print("\nFGSM攻击下的准确率比较:")
    for i, eps in enumerate(test_epsilons):
        print(f"Epsilon={eps:.2f}: 原始模型={orig_attack_accs[i]*100:.2f}%, "
              f"增强模型={robust_attack_accs[i]*100:.2f}%")
    
    print("\n在线对抗训练实验完成！")


if __name__ == "__main__":
    main()